#!/usr/bin/env python

# This script plots the distribution of the gop data.
#
# Inputs (given as arguments): video1_framestats.txt, video2_framestats.txt, etc...
# Output Plot #3: Distribution of GOP data size (for each file given)
#
#  To Do:
#	- Verify x-axis scaling

import numpy as np
import pylab as pl # Plotting
import os.path as osp # path
import sys # Handling command line arguments
from matplotlib.font_manager import FontProperties as FP # Font formatting in plots

# Helper function for analysing *_framestats.txt logs.

def new_frame(filep):
	line = '1'
	while not line.startswith('[/FRAME]'):
		line = filep.readline()
		if line.startswith('pkt_pos='):
			pos = int(line[line.find('=')+1:])
		if line.startswith('pict_type='):
			pict_type = line[line.find('=')+1:-1]
	return pos, pict_type

# Plotter function, is called with a reference to a *_framestats.txt file and a reference number.

def plot(filep,video_nb,tracker):

	##### DATA EXTRACTION #####

	pos_type = []

	line = '1'
	while line:
		line = f.readline()
		if line.startswith('[FRAME]'):
			position, pict_type = new_frame(f)
			pos_type.append((position, pict_type))

	# Prepare to build a list containing the data sizes of all GOP in video
	gopsize_list=[]
	gop_temp_size=0
	i_ref=0

	# Build gop_size_list 
	for i in range(len(pos_type)-1)	:
		if (str(pos_type[i+1][1]).lower()=='i'):
			gop_temp_size=pos_type[i][0]-pos_type[i_ref][0]
			gopsize_list.append(gop_temp_size)
			i_ref=i

	# Scale gop_size_list to kbytes - can be included in above loop
	for i in range(len(gopsize_list)): 	
		gopsize_list[i]=gopsize_list[i]/1000

	# Parameters for the histogram
#	histogram_bin_size = 2000 # OLD:10000 Carefully chosen
#	max_size=1 # since we have normalized gop sizes relative to total filesize
#	min_size=0

	# Parameters for the histogram
	histogram_bin_size = 150 # OLD: 50 Carefully chosen
	max_size=3000 # [kB] since we have normalized gop sizes relative to total filesize
	min_size=0

	# Check if we chose max_size properly
	if max_size<max(gopsize_list):
		print 'WARNING: We are missing some GOP in the Plot, increase max_size [kB]'

	total_size=pos_type[-1][0]/float(1000)

	# NB: We shall not normalize x anymore
	# Normalizing gop_size_list before making a histogram
	#	norm_gopsize_list=[]
	#	for i in range(len(gopsize_list)):
	#		norm_gopsize_list.append(gopsize_list[i]/float(total_size))

	# We are no longer normalizing x
	norm_gopsize_list=gopsize_list
	
	# make histogram of data
	hist_gopsize = pl.histogram(norm_gopsize_list, histogram_bin_size, (min_size, max_size))

	# Normalizing y data
	norm_hist_gopsize_y=[]
	nb_gops=np.sum(hist_gopsize[0])
	for i in range(len(hist_gopsize[0])):
		norm_hist_gopsize_y.append(hist_gopsize[0][i]/float(nb_gops))

	# Histogram data is already normalized - just calling it something else
	norm_hist_gopsize_x=hist_gopsize[1][:-1]

	# Prepare figure
	fig = pl.figure(figsize=[10,4]) # ugly way to set aspect ratio - works okay
	ax1 = fig.add_subplot(111)

	# Layout fix: Convert 0 to -1 in y_data which cleans plot
	for k in range(len(norm_hist_gopsize_y)):
		if norm_hist_gopsize_y[k]==0:
			norm_hist_gopsize_y[k]=-1

	# Layout fix: Numerate legend with letters: A,B,C (expect to only go to C) ...

	if video_nb<3:
		letter='A.'+str((tracker%3)+1)
	elif video_nb<6:
		letter='B.'+str((tracker%3)+1)
	else:
		letter='C.'+str((tracker%3)+1)
		
#	ax1.plot(norm_hist_gopsize_x, norm_hist_gopsize_y, linestyle='-',label='Video '+str(video_nb+1),drawstyle='steps-post')

	ax1.plot(norm_hist_gopsize_x, norm_hist_gopsize_y, linestyle='-',color='black',label='Video '+letter, drawstyle='steps-post')

	# Print useful information
	print 'Video '+str(video_nb+1)
	print 'Filesize '+str(total_size/1000.)+' MB'
	print 'Number of gops '+str(nb_gops)
	print 'Mean GOP size '+str(np.mean(gopsize_list))+'kB'
	print 'Mode GOP size '+str(norm_hist_gopsize_x[np.argmax(norm_hist_gopsize_y)])+'kB'
	print '\n'

	# Font fix
	fp = FP()
	fp.set_size('small')

 	# Plot annotation
	ax1.legend(prop=fp)
	ax1.set_xlabel('GOP size [kB]')
	ax1.set_ylabel('Part of all GOP [-]')
	ax1.grid('on')
	
	# Forcing similar axes across plots
	pl.ylim(0,0.35)
	pl.xlim(0,2000)

	# Print figure
	workingdir = osp.abspath('../../plotting/video_coding2')
	pl.savefig(workingdir+'/figs/dist_gop_'+str(video_nb+1)+'.eps')

if __name__ == '__main__':

	# Legend notation
	tracker=0

	# Run through all of the given gopstats.txt logfiles given as arguments
	for arg in range(len(sys.argv)-1):
		input_filename = sys.argv[arg+1]
		f = open(input_filename,'r')
		plot(f,arg,tracker)
		tracker+=1



